package com.ihai.spring.boot.sample.distributed.counter.util;

import com.ihai.spring.boot.sample.distributed.counter.manager.CounterManagerMBean;
import com.ihai.spring.boot.sample.distributed.counter.remote.Node;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.management.*;
import javax.management.remote.*;
import java.rmi.registry.LocateRegistry;
import java.rmi.server.ExportException;

/**
 * DESC
 *
 * @author 张驰
 * @date 2019/12/25
 */
public class JmxRemoteUtil {
    private static final Logger logger = LoggerFactory.getLogger(JmxRemoteUtil.class);
    private static JMXConnectorServer jmxConnectorServer;
    public static JMXConnector initConnection(Node node,String domain){
        try {
            JMXServiceURL serviceURL = new JMXServiceURL("service:jmx:rmi:///jndi/rmi://"+node.getIp()+":"+node.getPort()+"/"+domain);
            JMXConnector jmxConnector = JMXConnectorFactory.connect(serviceURL, null);
            return jmxConnector;
        } catch (Exception e){
            logger.error("get connection error.", e);
        }
        return null;
    }
    public static void startJmxServer(Node node, String domain){
        try {
            MBeanServer mBeanServer = MBeanServerFactory.createMBeanServer(domain);
            try{
                LocateRegistry.createRegistry(node.getPort());
            }catch (ExportException e){
                logger.warn("MBean Registry [" + node.getPort() + "] already created.");
            }
            JMXServiceURL url = new JMXServiceURL("service:jmx:rmi:///jndi/rmi://"+node.getIp()+":" + node.getPort() + "/"+domain);
            jmxConnectorServer = JMXConnectorServerFactory.newJMXConnectorServer(url, null, mBeanServer);
            jmxConnectorServer.start();
        } catch (Exception e){
            logger.error("start jmx server error.", e);
        }
    }

    public static void register(Object object, String objName){
        try {
            ObjectName objectName = new ObjectName(objName);
            jmxConnectorServer.getMBeanServer().registerMBean(object, objectName);
        } catch (Exception e){
            logger.error("register mbean error.", e);
        }
    }

    public static void unRegister(String objName){
        try {
            ObjectName objectName = new ObjectName(objName);
            jmxConnectorServer.getMBeanServer().unregisterMBean(objectName);
        } catch (Exception e){
            logger.error("register mbean error.", e);
        }
    }

    public static <T> T getMBean(Node node, String domain, String objName, Class<T> clazz){
        try {
            JMXConnector jmxConnector = initConnection(node,domain);
            if(jmxConnector != null){
                MBeanServerConnection connection = jmxConnector.getMBeanServerConnection();
                ObjectName objectName = new ObjectName(objName);
                if (!connection.isRegistered(objectName)) {
                    connection.createMBean(CounterManagerMBean.class.getName(), objectName, null, null);
                }
                return JMX.newMBeanProxy(connection, objectName, clazz);
            }
        }catch (Exception e){
            logger.error("get mbean failed.", e);
        }
        return null;
    }

}
